import random
from dashscope import Generation
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from http import HTTPStatus

app = APIRouter()


@app.websocket('/index')  # 因为当前接口写在chat分组应用中，所以websocket接口的访问地址：ws://127.0.0.1:8000/chat/index
async def index(websocket: WebSocket):
    """AI应用助理的聊天接口"""
    await websocket.accept()
    try:
        while True:
            prompt = await websocket.receive_text()
            messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},
                        {'role': 'user', 'content': ''}]
            response = Generation.call(model="qwen-turbo",
                                       prompt=prompt,
                                       api_key='sk-cd792be673964a5186fcded28d744a1c',
                                       messages=messages,
                                       # 设置随机数种子seed，如果没有设置，则随机数种子默认为1234
                                       seed=random.randint(1, 10000),
                                       # 将输出设置为"message"格式
                                       result_format='message'
                                       )
            if response.status_code == HTTPStatus.OK:
                print(response)
            else:
                print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
                    response.request_id, response.status_code,
                    response.code, response.message
                ))
            await websocket.send_text(response.output.choices[0].message.content.strip())
    except WebSocketDisconnect as e:
        pass
